import base64
import os
import time
from functools import wraps
from typing import Any, Callable, Type, Union

from basic.base_driver import DRIVER
from const import SCREEN_PATH


def exception_handing(obj: Union[Callable, Type]) -> Union[Callable, Type]:
    """
    异常处理装饰器，可以应用于类或函数
    
    当应用于类时：会自动为类中的所有测试方法添加异常处理
    当应用于函数时：直接为函数添加异常处理
    
    功能：
    1. 开始录屏
    2. 执行测试
    3. 异常时保存录屏并重新抛出异常
    """
    
    def method_decorator(func: Callable) -> Callable:
        """装饰单个方法的装饰器"""
        @wraps(func)
        def wrapper(*args, **kwargs):
            # 确保录屏目录存在
            os.makedirs(SCREEN_PATH, exist_ok=True)
            
            # 开始录屏
            try:
                DRIVER.start_recording_screen()
            except Exception:
                # 如果录屏失败，继续执行测试
                pass
            
            try:
                # 执行原始方法
                result = func(*args, **kwargs)
                return result
            except Exception as e:
                # 异常时保存录屏
                try:
                    result = DRIVER.stop_recording_screen()
                    if result:
                        screen = base64.b64decode(result)
                        timestamp = time.strftime("%Y%m%d_%H%M%S")
                        filename = f"{func.__name__}_{timestamp}.mp4"
                        filepath = os.path.join(SCREEN_PATH, filename)
                        with open(filepath, 'wb') as f:
                            f.write(screen)
                except Exception:
                    # 录屏处理失败时不影响异常抛出
                    pass
                
                # 重新抛出原始异常，让pytest正常处理
                raise e
            finally:
                # 确保录屏被停止
                try:
                    DRIVER.stop_recording_screen()
                except Exception:
                    pass
        
        return wrapper
    
    # 如果是类，装饰类中的所有测试方法
    if isinstance(obj, type):
        # 获取类中所有方法
        for attr_name in dir(obj):
            if not attr_name.startswith('_'):  # 忽略私有方法
                attr = getattr(obj, attr_name)
                if callable(attr) and attr_name.startswith('test_'):  # 只装饰测试方法
                    setattr(obj, attr_name, method_decorator(attr))
        return obj
    
    # 如果是函数，直接装饰函数
    elif callable(obj):
        return method_decorator(obj)
    
    else:
        raise TypeError("exception_handing 只能应用于类或函数")
